// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/coalesce_expr.h>

namespace kumo::pollux::exec {
    CoalesceExpr::CoalesceExpr(
        TypePtr type,
        std::vector<ExprPtr> &&inputs,
        bool inputsSupportFlatNoNullsFastPath)
        : SpecialForm(
            std::move(type),
            std::move(inputs),
            kCoalesce,
            inputsSupportFlatNoNullsFastPath,
            false /* trackCpuUsage */) {
        std::vector<TypePtr> inputTypes;
        inputTypes.reserve(inputs_.size());
        std::transform(
            inputs_.begin(),
            inputs_.end(),
            std::back_inserter(inputTypes),
            [](const ExprPtr &expr) { return expr->type(); });

        // Apply type checks.
        auto expectedType = resolveType(inputTypes);
        POLLUX_CHECK(
            *expectedType == *this->type(),
            "Coalesce expression type different than its inputs. Expected {} but got Actual {}.",
            expectedType->toString(),
            this->type()->toString());
    }

    void CoalesceExpr::evalSpecialForm(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        // Null positions to populate.
        exec::LocalSelectivityVector activeRowsHolder(context, rows.end());
        auto activeRows = activeRowsHolder.get();
        assert(activeRows); // for lint
        *activeRows = rows;

        // Fix finalSelection at "rows" unless already fixed.
        ScopedFinalSelectionSetter scopedFinalSelectionSetter(context, &rows);

        exec::LocalDecodedVector decodedVector(context);
        for (int i = 0; i < inputs_.size(); i++) {
            inputs_[i]->eval(*activeRows, context, result);

            if (!result->may_have_nulls()) {
                // No nulls left.
                return;
            }

            if (context.errors()) {
                context.deselectErrors(*activeRows);
            }

            decodedVector.get()->decode(*result, *activeRows);
            const uint64_t *rawNulls = decodedVector->nulls(activeRows);
            if (!rawNulls) {
                // No nulls left.
                return;
            }

            activeRows->deselectNonNulls(rawNulls, 0, activeRows->end());
            if (!activeRows->hasSelections()) {
                // No nulls left.
                return;
            }
        }
    }

    // static
    TypePtr CoalesceExpr::resolveType(const std::vector<TypePtr> &argTypes) {
        POLLUX_CHECK_GT(
            argTypes.size(),
            0,
            "COALESCE statements expect to receive at least 1 argument, but did not receive any.");
        for (auto i = 1; i < argTypes.size(); i++) {
            POLLUX_USER_CHECK(
                *argTypes[0] == *argTypes[i],
                "Inputs to coalesce must have the same type. Expected {}, but got {}.",
                argTypes[0]->toString(),
                argTypes[i]->toString());
        }

        return argTypes[0];
    }

    TypePtr CoalesceCallToSpecialForm::resolveType(
        const std::vector<TypePtr> &argTypes) {
        return CoalesceExpr::resolveType(argTypes);
    }

    ExprPtr CoalesceCallToSpecialForm::constructSpecialForm(
        const TypePtr &type,
        std::vector<ExprPtr> &&compiledChildren,
        bool /* trackCpuUsage */,
        const core::QueryConfig & /*config*/) {
        bool inputsSupportFlatNoNullsFastPath =
                Expr::allSupportFlatNoNullsFastPath(compiledChildren);
        return std::make_shared<CoalesceExpr>(
            type, std::move(compiledChildren), inputsSupportFlatNoNullsFastPath);
    }
} // namespace kumo::pollux::exec
